import concurrent
import multiprocessing
import os
import time
import concurrent.futures
import path
import librosa
import numpy as np
import onnxruntime as ort
import soundfile as sf
from tqdm import tqdm
from gxl_ai_utils.utils import utils_file
from gxl_ai_utils.thread.my_thread import MyThreadPool

available_providers = ort.get_available_providers()
print(available_providers)
# 输出: ['AzureExecutionProvider', 'CPUExecutionProvider']

SAMPLING_RATE = 16000
INPUT_LENGTH = 9.01


class ComputeScore:
    def __init__(self, primary_model_path, p808_model_path) -> None:
        self.onnx_sess = ort.InferenceSession(primary_model_path, providers=[available_providers[0]])
        self.p808_onnx_sess = ort.InferenceSession(p808_model_path,
                                                   providers=[available_providers[0]])

    def audio_melspec(self, audio, n_mels=120, frame_size=320, hop_length=160, sr=16000, to_db=True):
        mel_spec = librosa.feature.melspectrogram(y=audio, sr=sr, n_fft=frame_size + 1, hop_length=hop_length,
                                                  n_mels=n_mels)
        if to_db:
            mel_spec = (librosa.power_to_db(mel_spec, ref=np.max) + 40) / 40
        return mel_spec.T

    def get_polyfit_val(self, sig, bak, ovr, is_personalized_MOS):
        if is_personalized_MOS:
            p_ovr = np.poly1d([-0.00533021, 0.005101, 1.18058466, -0.11236046])
            p_sig = np.poly1d([-0.01019296, 0.02751166, 1.19576786, -0.24348726])
            p_bak = np.poly1d([-0.04976499, 0.44276479, -0.1644611, 0.96883132])
        else:
            p_ovr = np.poly1d([-0.06766283, 1.11546468, 0.04602535])
            p_sig = np.poly1d([-0.08397278, 1.22083953, 0.0052439])
            p_bak = np.poly1d([-0.13166888, 1.60915514, -0.39604546])

        sig_poly = p_sig(sig)
        bak_poly = p_bak(bak)
        ovr_poly = p_ovr(ovr)

        return sig_poly, bak_poly, ovr_poly

    def __call__(self, fpath, sampling_rate, is_personalized_MOS=False):
        aud, input_fs = sf.read(fpath)
        fs = sampling_rate
        if input_fs != fs:
            audio = librosa.resample(aud, input_fs, fs)
        else:
            audio = aud
        actual_audio_len = len(audio)
        len_samples = int(INPUT_LENGTH * fs)
        while len(audio) < len_samples:
            audio = np.append(audio, audio)

        num_hops = int(np.floor(len(audio) / fs) - INPUT_LENGTH) + 1
        hop_len_samples = fs
        predicted_mos_sig_seg_raw = []
        predicted_mos_bak_seg_raw = []
        predicted_mos_ovr_seg_raw = []
        predicted_mos_sig_seg = []
        predicted_mos_bak_seg = []
        predicted_mos_ovr_seg = []
        predicted_p808_mos = []

        for idx in range(num_hops):
            audio_seg = audio[int(idx * hop_len_samples): int((idx + INPUT_LENGTH) * hop_len_samples)]
            if len(audio_seg) < len_samples:
                continue

            input_features = np.array(audio_seg).astype('float32')[np.newaxis, :]
            p808_input_features = np.array(self.audio_melspec(audio=audio_seg[:-160])).astype('float32')[np.newaxis, :,
                                  :]
            oi = {'input_1': input_features}
            p808_oi = {'input_1': p808_input_features}
            p808_mos = self.p808_onnx_sess.run(None, p808_oi)[0][0][0]
            mos_sig_raw, mos_bak_raw, mos_ovr_raw = self.onnx_sess.run(None, oi)[0][0]
            mos_sig, mos_bak, mos_ovr = self.get_polyfit_val(mos_sig_raw, mos_bak_raw, mos_ovr_raw, is_personalized_MOS)
            predicted_mos_sig_seg_raw.append(mos_sig_raw)
            predicted_mos_bak_seg_raw.append(mos_bak_raw)
            predicted_mos_ovr_seg_raw.append(mos_ovr_raw)
            predicted_mos_sig_seg.append(mos_sig)
            predicted_mos_bak_seg.append(mos_bak)
            predicted_mos_ovr_seg.append(mos_ovr)
            predicted_p808_mos.append(p808_mos)

        clip_dict = {'filename': fpath, 'len_in_sec': actual_audio_len / fs, 'sr': fs}
        clip_dict['num_hops'] = num_hops
        clip_dict['OVRL_raw'] = np.mean(predicted_mos_ovr_seg_raw)
        clip_dict['SIG_raw'] = np.mean(predicted_mos_sig_seg_raw)
        clip_dict['BAK_raw'] = np.mean(predicted_mos_bak_seg_raw)
        clip_dict['OVRL'] = np.mean(predicted_mos_ovr_seg)
        clip_dict['SIG'] = np.mean(predicted_mos_sig_seg)
        clip_dict['BAK'] = np.mean(predicted_mos_bak_seg)
        clip_dict['P808_MOS'] = np.mean(predicted_p808_mos)
        return clip_dict['P808_MOS']


def get_mos_for_scp(wav_scp_input_path: str, mos_scp_output_path: str):
    """
    Main function, 对数据进行mos打分
    """
    utils_file.makedir_for_file(mos_scp_output_path)
    dic = utils_file.load_dic_from_scp(wav_scp_input_path)
    res_dic = {}
    p808_model_path = os.path.join('DNSMOS', 'model_v8.onnx')
    primary_model_path = os.path.join('pDNSMOS', 'sig_bak_ovr.onnx')
    compute_score = ComputeScore(primary_model_path, p808_model_path)
    print('开始得到得到wav.scp文件对应音频的p808mos得分, 共有{}条数据'.format(len(dic)))
    for k, v in tqdm(dic.items(), total=len(dic)):
        res_dic[k] = compute_score(v, sampling_rate=SAMPLING_RATE)
    utils_file.write_dic_to_scp(res_dic, mos_scp_output_path)


def get_mos_for_scp_multi_thread(wav_scp_input_path: str, mos_scp_output_path: str):
    """
    Main function, 对数据进行mos打分
    """
    num_thread = 5
    utils_file.makedir_for_file(mos_scp_output_path)
    dic = utils_file.load_dic_from_scp(wav_scp_input_path)
    res_dic = {}
    p808_model_path = os.path.join('DNSMOS', 'model_v8.onnx')
    primary_model_path = os.path.join('pDNSMOS', 'sig_bak_ovr.onnx')
    compute_score = ComputeScore(primary_model_path, p808_model_path)
    print('开始得到得到wav.scp文件对应音频的p808mos得分, 共有{}条数据'.format(len(dic)))

    def hande_dic(dict_o: dict):
        now_dic = dict_o
        for k, v in tqdm(now_dic.items(), total=len(now_dic)):
            res_dic[k] = compute_score(v, sampling_rate=SAMPLING_RATE)

    item_list = list(dic.items())
    my_thread_pool = MyThreadPool()
    split_num = len(item_list) // num_thread
    for i in range(0, len(item_list), split_num):
        now_dic = {k: v for k, v in item_list[i:i + split_num]}
        my_thread_pool.add_thread(hande_dic, [now_dic])
    my_thread_pool.start()

    utils_file.write_dic_to_scp(res_dic, mos_scp_output_path)


if __name__ == '__main__':
    """"""
    input_scp = './output/wav.scp'
    output_scp = './output/p808_mos.scp'
    print(f'总共需要处理{len(utils_file.load_dic_from_scp("./output/wav.scp"))}条数据')
    now = time.time()
    get_mos_for_scp_multi_thread(input_scp, output_scp)
    end = time.time()
    print(f'耗时{end - now}s')
